#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

import heapq
# import math
from math import *


class Node:
    # __slots__ = ['']

    def __init__(self, freq, ele=None, left=None, right=None):
        '''
        ele可能为空, 非叶子节点
        '''
        freq = float(freq)

        self.coding = None
        self.freq = freq
        self.ele = ele
        self.left = left
        self.right = right

    def __lt__(self, other):
        # 用于堆/优先级队列排序: 插入或者弹出
        return self.freq < other.freq

    def show(self):
        print('[%s] : [%s] [%s]' % (self.ele, self.freq, self.coding))

    def __add__(self, other):
        node = Node(self.freq + other.freq, None, self, other)
        return node


class HuffmanCode:
    __slots__ = ['heap_l', 'huff_tree', 'weight_map', 'coding_map']

    def __init__(self, nodes_list):
        nodes = [Node(freq, ele) for ele, freq in nodes_list]
        heapq.heapify(nodes)

        self.heap_l = nodes
        self.coding_map = {node.ele:node for node in nodes}
        # print(self.coding_map)
        self.huff_tree = self.build_huffman()
        self._record_coding()
        self.show_coding()

        self.weight()

    def build_huffman(self):
        ''' 建立哈夫曼树, 用于后续编码 '''
        n1, n2 = self._pop_2()
        while n1 and n2:
            # print('===================  =====================')
            # n1.show()
            # n2.show()
            new_node = n2 + n1      # 频率低的放在右边
            # new_node = n1 + n2      # 频率低的放在左边
            heapq.heappush(self.heap_l, new_node)
            n1, n2 = self._pop_2()

        # n1.show()
        # print(n1, n2)
        return n1

    def _pop_2(self):
        n1, n2 = None, None
        try:
            n1 = heapq.heappop(self.heap_l)
            n2 = heapq.heappop(self.heap_l)
        except:
            pass

        return n1, n2

    def _record_coding(self):
        assert self.huff_tree
        print('===================  =====================')
        self._travel(self.huff_tree, '')

    def _travel(self, node, prefix):
        if node.ele:
            # 带有元素的节点, 一定是叶子节点: 其路径为编码
            self.coding_map[node.ele].coding = prefix
            return

        # 非叶子节点, 一定不是元素节点, 并且一定有左右子节点
        self._travel(node.left,  prefix+'0')
        self._travel(node.right, prefix+'1')

    def show_coding(self):
        # for ele, node in self.coding_map.items():
        for node in self.coding_map.values():
            node.show()

    def weight(self):
        huffman_sum = sum([node.freq * len(node.coding) for node in self.coding_map.values()])
        print('huffman weight sum [%s]' % huffman_sum)

        l = len(self.coding_map)
        stad_len = ceil(log(l, 2))
        stad_sum = sum( [stad_len * node.freq for node in self.coding_map.values()] )
        print('stad len [%s]' % stad_len)
        print('stad weight sum [%s]' % stad_sum)


def main():
    print("start main")

    nodes_list = [
        (56, 'e'),
        (23, 'k'),
        (32, 'j'),
        (5,  'z')
    ]

    nodes_list = zip(
        ['a', 'b', 'c', 'd', 'e'],
        [5, 4, 3, 2, 1],
    )

    h = HuffmanCode(nodes_list)


if __name__ == "__main__":
    main()
